package com.wenjie.common.reRequest;

import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

public class CaptureResponseWrapper extends HttpServletResponseWrapper {
    private ByteArrayOutputStream baos = new ByteArrayOutputStream();
    private CapturingServletOutputStream capturingStream = new CapturingServletOutputStream(baos);

    public CaptureResponseWrapper(HttpServletResponse response) {
        super(response);
    }

    @Override
    public javax.servlet.ServletOutputStream getOutputStream() throws IOException {
        return capturingStream;
    }

    public String getCapturedResponse() {
        return baos.toString();
    }

    private class CapturingServletOutputStream extends javax.servlet.ServletOutputStream {
        private ByteArrayOutputStream baos;

        public CapturingServletOutputStream(ByteArrayOutputStream baos) {
            this.baos = baos;
        }

        @Override
        public void write(int b) throws IOException {
            baos.write(b);
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setWriteListener(WriteListener writeListener) {

        }
    }
}
